import pymongo
from pymongo import MongoClient
from pymongo.errors import ConnectionFailure


class MongoDBUtils:
    def __init__(self, uri, db_name, max_connections=10):
        """
        初始化MongoDB连接池工具类。

        :param uri: MongoDB连接URI
        :param db_name: 要连接的数据库名称
        :param max_connections: 最大连接数，默认为10
        """
        self.uri = uri
        self.db_name = db_name
        self.max_connections = max_connections
        self.client = None
        self.lock = pymongo.threading.Lock()  # 线程锁
        self.connect()

    def connect(self):
        """
        创建MongoDB连接。
        """
        try:
            with self.lock:
                if not self.client:
                    self.client = MongoClient(
                        self.uri, maxPoolSize=self.max_connections
                    )
        except ConnectionFailure as e:
            print(f"Could not connect to MongoDB: {e}")

    def get_database(self):
        """
        获取数据库对象。

        :return: 数据库对象
        """
        self.connect()
        return self.client[self.db_name]

    def get_collection(self, collection_name):
        """
        获取集合对象。

        :param collection_name: 集合名称
        :return: 集合对象
        """
        db = self.get_database()
        return db[collection_name]

    def insert_one(self, collection_name, document):
        """
        向指定集合插入单个文档。

        :param collection_name: 集合名称
        :param document: 要插入的文档
        :return: 插入结果
        """
        collection = self.get_collection(collection_name)
        return collection.insert_one(document)

    def insert_many(self, collection_name, documents):
        """
        向指定集合插入多个文档。

        :param collection_name: 集合名称
        :param documents: 要插入的文档列表
        :return: 插入结果
        """
        collection = self.get_collection(collection_name)
        return collection.insert_many(documents)

    def find(self, collection_name, query):
        """
        查询指定集合中的文档。

        :param collection_name: 集合名称
        :param query: 查询条件
        :return: 查询结果
        """
        collection = self.get_collection(collection_name)
        return collection.find(query)

    def update_one(self, collection_name, filter, update):
        """
        更新指定集合中的单个文档。

        :param collection_name: 集合名称
        :param filter: 过滤条件
        :param update: 更新操作
        :return: 更新结果
        """
        collection = self.get_collection(collection_name)
        return collection.update_one(filter, update)

    def delete_one(self, collection_name, filter):
        """
        删除指定集合中的单个文档。

        :param collection_name: 集合名称
        :param filter: 过滤条件
        :return: 删除结果
        """
        collection = self.get_collection(collection_name)
        return collection.delete_one(filter)

    def close(self):
        """
        关闭MongoDB连接。
        """
        if self.client:
            self.client.close()


# 使用示例
if __name__ == "__main__":
    uri = "mongodb://localhost:27017/"
    db_name = "test_db"
    # 创建连接池实例
    mongo_utils = MongoDBUtils(uri, db_name)
    # 获取数据库
    db = mongo_utils.get_database()
    # 插入数据示例
    sample_data = {"name": "Alice", "age": 30}
    result = mongo_utils.insert_one("users", sample_data)

    print(f"Inserted document with _id: {result.inserted_id}")
    # 查询数据示例
    user = next(mongo_utils.find("users", {"name": "Alice"}))
    print(user)
    # 关闭连接
    mongo_utils.close()
